#include "sample.h"
#include <cuda.h>


namespace uzu
{
    __global__ void VecAddGpuKernel(float* A, float* B, float* C, int N)
    {
        int i = blockDim.x * blockIdx.x + threadIdx.x;
        if (i < N)
            C[i] = A[i] + B[i];
    }

    void VecAddGpu(float* A, float* B, float* C, int N)
    {
        // allocate device memory
        int buffer_size = N * sizeof(float);
        float* device_A;
        float* device_B;
        float* device_C;

        cudaMalloc(&device_A, buffer_size);
        cudaMalloc(&device_B, buffer_size);
        cudaMalloc(&device_C, buffer_size);

        cudaMemcpy(device_A, A, buffer_size, cudaMemcpyHostToDevice);
        cudaMemcpy(device_B, B, buffer_size, cudaMemcpyHostToDevice);

        // launch kernel
        int block_size = 256;
        int num_blocks = (N + block_size - 1) / block_size;
        VecAddGpuKernel<<<num_blocks, block_size>>>(device_A, device_B, device_C, N);
        // copy data back and free device memory
        cudaMemcpy(C, device_C, buffer_size, cudaMemcpyDeviceToHost);
        cudaFree(device_A);
        cudaFree(device_B);
        cudaFree(device_C);
    }
}
